Greptile Summary

This PR fixes a bug in get_attention_backend where FlashAttention 2 was incorrectly disabled for head_dim > 192 on architectures not in an exact-match allowlist (sm80, sm90, sm100, sm120), excluding valid devices like sm103, sm89, and sm86. The fix replaces the allowlist with the simpler head_dim_qk > 256 or head_dim_qk % 8 != 0 condition, correctly aligning with flash-attn's own >= sm80 support guarantee — which is already enforced earlier in the function at the compute-capability filter (lines 448–451).

Confidence Score: 5/5

Safe to merge — minimal, targeted bug fix with correct logic and no regressions introduced.

The change removes a single, clearly erroneous allowlist condition. The new condition is logically equivalent to flash-attn's own gate given the earlier < sm80 guard already disables FA2 before this point. The dead-code concern flagged in the previous review thread is fully resolved by removing the branch entirely. No new logic is added, and the log message is updated consistently.

No files require special attention.

Important Files Changed

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[get_attention_backend called] --> B{device_compute_capability < sm80?}
    B -- Yes --> C[use_flash_attention_2 = False]
    B -- No --> D{use_flash_attention_2 AND FA2 installed?}
    D -- No --> G[Skip FA2 head_dim check]
    D -- Yes --> E{head_dim_qk > 256\nOR head_dim_qk % 8 != 0?}
    E -- Yes --> F[use_flash_attention_2 = False\nlog debug message]
    E -- No --> H[FA2 remains enabled]
    C --> I[Continue backend selection]
    F --> I
    G --> I
    H --> I

    style C fill:#f88,stroke:#c00
    style F fill:#f88,stroke:#c00
    style H fill:#8f8,stroke:#090

_{Reviews (3): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

/te-ci L0

/te-ci pytorch